/**
* @file model.cpp
*
* Copyright (c) Huawei Technologies Co., Ltd. 2023. All Rights reserved.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*/
#include "acl/acl_mdl.h"
#include <securec.h>
#include "acl/acl_base.h"
#include "framework/executor_c/ge_executor.h"
#include "model_desc_internal.h"
#include "log_inner.h"
#include "toolchain/slog.h"
#include "vector.h"
#include "runtime/rt.h"
#include "model_config.h"
static aclError ModelLoadFromFileWithMem(const char *const modelPath, uint32_t *const modelId, void *const weightPtr,
                                         const size_t weightSize, const int32_t priority) {
  uint32_t id = 0U;
  ModelData data;
  ACL_LOG_INFO("call ge interface executor.LoadDataFromFile, weightSize[%zu]", weightSize);
  Status ret = LoadDataFromFile(modelPath, &data);
  if (ret != SUCCESS) {
    ACL_LOG_CALL_ERROR("[Model][FromFile]load model from file[%s] failed, ge result[%u]", modelPath, ret);
    return ret;
  }
  data.priority = priority;
  ret = LoadModelFromData(&id, &data, weightPtr, weightSize);
  if (ret != SUCCESS) {
      FreeModelData(&data);
      ACL_LOG_CALL_ERROR("[Model][FromData]load model from data failed, ge result[%u]", ret);
      return ret;
  }

  *modelId = id;
  FreeModelData(&data);
  ACL_LOG_INFO("successfully execute ModelLoadFromFileWithMem, weightSize[%zu], modelId[%u]",
      weightSize, *modelId);
  return ACL_SUCCESS;
}

static aclError ModelLoadFromMemWithMem(const void *const model, const size_t modelSize, uint32_t *const modelId,
                                        void *const weightPtr, const size_t weightSize, const int32_t priority) {
  ACL_LOG_INFO("start to execute ModelLoadFromMemWithMem, weightSize[%zu], priority[%d]",
      weightSize, priority);
  if (modelSize == 0U) {
    ACL_LOG_INNER_ERROR("[Check][ModelSize]modelSize[%zu] is invalid, should not be zero", modelSize);
    return ACL_ERROR_INVALID_PARAM;
  }

  uint32_t id = 0U;
  ModelData modelData;
  modelData.modelData = (void *)model;
  modelData.modelLen = (uint64_t)modelSize;
  modelData.priority = priority;
  ACL_LOG_INFO("call ge interface executor.LoadModelFromData, modelSize[%zu], weightSize[%zu]",
      modelSize, weightSize);
  Status ret = LoadModelFromData(&id, &modelData, weightPtr, weightSize);
  if (ret != SUCCESS) {
      ACL_LOG_CALL_ERROR("[Model][FromData]load model from data failed, ge result[%u]", ret);
      return ret;
  }

  *modelId = id;
  ACL_LOG_INFO("successfully execute ModelLoadFromMemWithMem, modelSize[%zu], weightSize[%zu], modelId[%u]",
      modelSize, weightSize, *modelId);
  return ACL_SUCCESS;
}

aclError SetExecHandle(const aclmdlExecConfigHandle *handle, aclrtStream stream, ExecHandleDesc *execDesc) {
  void *workPtr = NULL;
  size_t workSize = 0U;
  if (handle != NULL) {
    if (handle->workPtr != NULL) {
      workPtr = handle->workPtr;
      workSize = handle->workSize;
    }
  }
  if ((workPtr == NULL) && (stream != NULL)) {
    (void)rtStreamGetWorkspace(stream, &workPtr, &workSize);
  }
  if (workPtr == NULL) {
    ACL_LOG_ERROR("stream and exeHandle is not set workSpacePtr.");
    return ACL_ERROR_INVALID_PARAM;
  }
  execDesc->stream = stream;
  execDesc->workPtr = workPtr;
  execDesc->workSize = workSize;
  if (handle != NULL) {
    execDesc->mpamId = handle->mpamId;
    execDesc->aicQos = handle->aicQos;
    execDesc->aicOst = handle->aicOst;
    execDesc->mecTimeThreshHold = handle->mecTimeThreshHold;
  }
  return ACL_SUCCESS;
}

static aclError ModelExecute(const uint32_t modelId, const aclmdlDataset *const input, aclmdlDataset *const output,
                             const bool sync, aclrtStream stream, const aclmdlExecConfigHandle *handle) {
  ACL_LOG_INFO("start to execute ModelExecute, modelId[%u]", modelId);
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(input);
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(output);
  ExecHandleDesc execDesc = {0};
  aclError err = SetExecHandle(handle, stream, &execDesc);
  if (err != ACL_SUCCESS) {
    ACL_LOG_ERROR("para set error.");
    return err;
  }

  ACL_LOG_INFO("call ge interface executor.ExecModel, modelId[%u], asyncMode[%d]", modelId, sync);
  const Status ret = ExecModel(modelId, &execDesc, sync, (InputData *)input, (OutputData *)output);
  if (ret != SUCCESS) {
    ACL_LOG_CALL_ERROR("[Exec][Model]Execute model failed, ge result[%u], modelId[%u]", ret, modelId);
    return ret;
  }
  ACL_LOG_INFO("successfully execute ModelExecute, modelId[%u]", modelId);
  return ACL_SUCCESS;
}

aclError aclmdlLoadWithConfig(const aclmdlConfigHandle *handle, uint32_t *modelId) {
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(handle);
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(modelId);
  if (!CheckMdlConfigHandle(handle)) {
    ACL_LOG_INNER_ERROR("[Check][ConfigHandle]model config is invalid because some param may not be set");
    return ACL_ERROR_INVALID_PARAM;
  }
  aclError ret;
  switch (handle->mdlLoadType) {
    case ACL_MDL_LOAD_FROM_FILE: {
      ret = ModelLoadFromFileWithMem(handle->loadPath, modelId, NULL, 0U, handle->priority);
      break;
    }
    case ACL_MDL_LOAD_FROM_FILE_WITH_MEM: {
      ret = ModelLoadFromFileWithMem(handle->loadPath, modelId, handle->weightPtr, handle->weightSize,
                                     handle->priority);
      break;
    }
    case ACL_MDL_LOAD_FROM_MEM: {
      ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(handle->mdlAddr);
      ret = ModelLoadFromMemWithMem(handle->mdlAddr, handle->mdlSize, modelId, NULL, 0U, handle->priority);
      break;
    }
    case ACL_MDL_LOAD_FROM_MEM_WITH_MEM: {
      ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(handle->mdlAddr);
      ret = ModelLoadFromMemWithMem(handle->mdlAddr, handle->mdlSize, modelId, handle->weightPtr,
                                    handle->weightSize, handle->priority);
      break;
    }
    default:
      ACL_LOG_INNER_ERROR("[Load][Model]model load type[%zu] is invalid, it should be in [%d, %d]",
          handle->mdlLoadType, ACL_MDL_LOAD_FROM_FILE, ACL_MDL_LOAD_FROM_MEM_WITH_MEM);
      return ACL_ERROR_INVALID_PARAM;
  }
  if (ret != ACL_SUCCESS) {
    ACL_LOG_INNER_ERROR("[Load][Model]model load type[%zu] failed.", handle->mdlLoadType);
    return ret;
  }
  ACL_LOG_INFO("success execute aclmdlLoadWithConfig, modeId[%u].", *modelId);
  return ACL_SUCCESS;
}


aclError aclmdlExecuteV2(uint32_t modelId, const aclmdlDataset *input, aclmdlDataset *output, aclrtStream stream,
                         const aclmdlExecConfigHandle *handle) {
  ACL_LOG_INFO("start to execute aclmdlExecuteV2, modelId[%u]", modelId);
  aclError ret = ModelExecute(modelId, input, output, true, stream, handle);
  if (ret == ACL_SUCCESS) {
    ACL_LOG_INFO("aclmdlExecuteV2 success, modelId[%u]", modelId);
  } else {
    ACL_LOG_INNER_ERROR("[Exec][Model]modelId[%u] execute failed, result[%d]", modelId, ret);
  }
  return ret;
}


aclError aclmdlExecuteAsyncV2(uint32_t modelId, const aclmdlDataset *input, aclmdlDataset *output, aclrtStream stream,
                              const aclmdlExecConfigHandle *handle) {
  ACL_LOG_INFO("start to execute aclmdlExecuteAsyncV2, modelId[%u]", modelId);
  aclError ret = ModelExecute(modelId, input, output, false, stream, handle);
  if (ret == ACL_SUCCESS) {
    ACL_LOG_INFO("aclmdlExecuteAsyncV2 success, modelId[%u]", modelId);
  } else {
    ACL_LOG_INNER_ERROR("[Exec][Model]modelId[%u] execute failed, result[%d]", modelId, ret);
  }
  return ret;
}

aclError aclmdlUnload(uint32_t modelId) {
  ACL_LOG_INFO("start to execute ACL_ModelUnload, modelId[%u]", modelId);
  const Status ret = UnloadModel(modelId);
  if (ret != SUCCESS) {
    ACL_LOG_CALL_ERROR("[Unload][Model]model unload failed, ge result[%u], modelId[%u]", ret, modelId);
    return ret;
  }
  ACL_LOG_INFO("aclmdlUnload success, modelId[%u]", modelId);
  return ACL_SUCCESS;
}


aclError aclmdlQuerySize(const char *fileName, size_t *workSize, size_t *weightSize) {
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(fileName);
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(workSize);
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(weightSize);
  const Status ret = GetMemAndWeightSize(fileName, workSize, weightSize);
  if (ret != SUCCESS) {
    ACL_LOG_CALL_ERROR("[Get][MemAndWeightSize]query size failed, ge result[%u]", ret);
    return ret;
  }
  ACL_LOG_INFO("success to get size from file[%s], work size[%zu] bytes, weight size[%zu] bytes",
      fileName, *workSize, *weightSize);
  return ACL_SUCCESS;
}

aclmdlDataset *aclmdlCreateDataset() {
  DataSet *dataset = (DataSet *)malloc(sizeof(DataSet));
  if (dataset == NULL) {
    ACL_LOG_ERROR("create data set failed.");
    return NULL;
  }
  InitVector(&dataset->blobs, sizeof(DataBlob));
  dataset->io_addr = NULL;
  dataset->ioa_size = 0U;
  return (aclmdlDataset *)dataset;
}

aclError aclmdlDestroyDataset(const aclmdlDataset *dataset) {
  ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(dataset);
  DataSet *geDataSet = (DataSet *)dataset;
  DeInitVector(&geDataSet->blobs);
  if (geDataSet->io_addr != NULL) {
    (void)rtFree(geDataSet->io_addr);
  }
  geDataSet->ioa_size = 0U;
  geDataSet->io_addr = NULL;
  free(geDataSet);
  return ACL_SUCCESS;
}

aclError aclmdlAddDatasetBuffer(aclmdlDataset *dataset, aclDataBuffer *dataBuffer) {
  if ((dataset == NULL) || (dataBuffer == NULL)) {
    ACL_LOG_ERROR("invalid input args.");
    return ACL_ERROR;
  }
  DataBuffer *geDatabuffer = (DataBuffer *)dataBuffer;
  DataSet *geDataSet = (DataSet *)dataset;
  DataBlob blob;
  blob.dataBuffer = geDatabuffer;
  if (EmplaceBackVector(&geDataSet->blobs, &blob) == NULL) {
    return ACL_ERROR;
  }
  return ACL_SUCCESS;
}

size_t aclmdlGetDatasetNumBuffers(const aclmdlDataset *dataset) {
  DataSet *geDataSet = (DataSet *)dataset;
  if (geDataSet == NULL) {
    ACL_LOG_ERROR("invalid input args.");
    REPORT_INPUT_ERROR(INVALID_PULL_POINTER_MSG, ARRAY("param"), ARRAY("dataset"));
    return 0;
  }
  return VectorSize(&geDataSet->blobs);
}

aclDataBuffer *aclmdlGetDatasetBuffer(const aclmdlDataset *dataset, size_t index) {
  DataSet *geDataSet = (DataSet *)dataset;
  if (geDataSet == NULL) {
    ACL_LOG_ERROR("invalid input args.");
    return NULL;
  }
  DataBlob *blob = VectorAt(&geDataSet->blobs, index);
  if (blob == NULL) {
    ACL_LOG_ERROR("data set has no data.");
    return NULL;
  }
  return (aclDataBuffer *)blob->dataBuffer;
}